import pandas
from vquant.store import Store


class DataFeed(Store):
    fields = dict(
        datetime=str, symbol=str, interval=str, open=float, high=float,
        low=float, close=float, volume=float
    )

    M1, M5, M15, M30, H, D, W, M = (
        '1min', '5min', '15min', '30min', 'H', 'D', 'W', 'M'
    )
    Intervals = (M1, M5, M15, M30, H, D, W, M)

    def __init__(self, cerebro):
        super(DataFeed, self).__init__()
        self.cerebro = cerebro
        self.datetime = pandas.Timestamp.now()
        self.deadline = pandas.Timestamp.now()

    @staticmethod
    def resample(dataframe: pandas.DataFrame, interval):
        dataframe.set_index(dataframe['datetime'], inplace=True)
        return dataframe.resample(interval).agg(
            dict(
                datetime='last', open='first', high='max',
                low='min', close='last', volume='sum'
            )
        ).dropna()

    def append(self, dataframes: [pandas.DataFrame], unit=None):
        self.dataframe = pandas.concat(dataframes, ignore_index=True)
        self.dataframe['datetime'] = pandas.to_datetime(
            self.dataframe['datetime'], unit=unit
        )

    def start(self):
        fragment = self.dataframe.loc[
            self.dataframe['datetime'].ge(self.datetime) &
            self.dataframe['datetime'].le(self.deadline)
        ]
        fragment = fragment.sort_values(by='datetime')
        fragment = fragment['datetime'].unique()
        for datetime in fragment:
            self.datetime = datetime
            dataframe = self.query(datetime=dict(eq=self.datetime))
            dataframe = dataframe.set_index(dataframe['symbol'])
            self.cerebro.on_tick(dataframe)


__all__ = [DataFeed]
